/**
 * 
 */
package totalFileSize;

import java.io.File;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author mingfan
 *
 * Calculate total file size of a given path
 * Use class CountDownLatch to synchronize threads
 */
public class TotalFileSizeWithLatch {
	
	// Total file size of the given path
	private final AtomicLong totalFileSize = new AtomicLong(0);
	
	// Number of sub-directories that need to be processed
	private final AtomicLong pendingFileVisits = new AtomicLong(0);
	
	// Thread pool
	private final ExecutorService threadPool;
	
	//
	private final CountDownLatch latch = new CountDownLatch(1);
	
	public TotalFileSizeWithLatch(final int threadPoolSize) {
		threadPool = Executors.newFixedThreadPool(threadPoolSize);
	}
	
	/**
	 * @param targetPath
	 */
	private void calculateTotalFileSize(final File targetPath) {

		long totalSize = 0;
		if (targetPath.isFile()) {
			totalSize = targetPath.length();
		} else {
			final File[] files = targetPath.listFiles();
			if (files != null) {
				for (final File oneFile : files) {
					if (oneFile.isFile()) {
						totalSize += oneFile.length();
					} else {
						pendingFileVisits.incrementAndGet();
						threadPool.execute(new Runnable() {

							@Override
							public void run() {
								calculateTotalFileSize(oneFile);
							}

						});
					}
				}
			}
		}
		totalFileSize.addAndGet(totalSize);
		if (pendingFileVisits.decrementAndGet() == 0) {
			latch.countDown();
		}
	}
	
	public long getTotalFileSize(final File targetPath) throws InterruptedException {
		pendingFileVisits.incrementAndGet();
		try {
			calculateTotalFileSize(targetPath);
			latch.await(30, TimeUnit.SECONDS);
//			latch.await();
			return totalFileSize.longValue();
		} finally {
			threadPool.shutdown();
		}
	}
	
	public static void main(String[] args) throws InterruptedException {
		
		if (args.length == 0) {
			System.out.println("No parameter given");
		} else {
			final String targetPath = args[0];
			System.out.println("Start calculate total file size of path "
					+ targetPath);
			final long start = System.nanoTime();
			final long totalSize = (new TotalFileSizeWithLatch(8))
					.getTotalFileSize(new File(targetPath));
			final long end = System.nanoTime();
			System.out.printf("Total file size of path %s is %d.%n",
					targetPath, totalSize);
			System.out.println("Time used is " + (end - start) / 1.0e9);
		}
	}

}
